﻿#include "ComDb.h"
#include "C0C_Def.h"

#include "params.h"
#include "Utils.h"
#include "PortNum.h"
#include "EnumDevices.h"

#include <stdio.h>
#include <RegStr.h>
#include <msports.h>
#include <Windows.h>
#include <QDebug>

#define TEXT_PREF

static const char comDbLocalKey[] = REGSTR_PATH_SERVICES "\\" C0C_SERVICE "\\COM Name Arbiter";
static const char comDbLocalName[] = "ComDB";

static WORD name2num(const char *pPortName)
{
    int num;

    if ((pPortName[0] != 'C' && pPortName[0] != 'c') ||
        (pPortName[1] != 'O' && pPortName[1] != 'o') ||
        (pPortName[2] != 'M' && pPortName[2] != 'm') ||
        pPortName[3] == '0' ||
        !Utils::StrToInt(pPortName + 3, &num) ||
        num <= 0 ||
        num > COMDB_MAX_PORTS_ARBITRATED)
    {
        return 0;
    }

    return (WORD)num;
}

static bool LoadComDb(PortBusyMask &comDb)
{
    comDb.clear();

    int res;

    do {
        res = IDCONTINUE;

        HCOMDB  hComDB;
        LONG err;

        err = ComDBOpen(&hComDB);

        if (err != ERROR_SUCCESS) {
            qDebug() << "ComDBOpen() failed with error:" << err;
            res = IDCANCEL;
            continue;
        }

        DWORD maxPortsReported;

        err = ComDBGetCurrentPortUsage(hComDB, NULL, 0, CDB_REPORT_BITS, &maxPortsReported);

        if (err != ERROR_SUCCESS) {
            ComDBClose(hComDB);
            qDebug() << "ComDBGetCurrentPortUsage() failed with error:" << err;
            res = IDCANCEL;
            continue;
        }

        DWORD bufSize = (maxPortsReported + 7)/8;
        BYTE *pBuf = (BYTE *)LocalAlloc(LPTR, bufSize);

        if (!pBuf) {
            ComDBClose(hComDB);

            qDebug() << "LocalAlloc(" << (unsigned long)bufSize << ") failed with ERROR_NOT_ENOUGH_MEMORY";
            res = IDCANCEL;
            continue;
        }

        err = ComDBGetCurrentPortUsage(hComDB, pBuf, bufSize, CDB_REPORT_BITS, &maxPortsReported);
        ComDBClose(hComDB);

        if (err != ERROR_SUCCESS) {
            LocalFree(pBuf);
            qDebug() << "ComDBGetCurrentPortUsage() failed with error:" << err;
            res = IDCANCEL;
            continue;
        }

        for (DWORD num = 0 ; num < maxPortsReported ; num++) {
            if (((pBuf[num/8] >> (num%8)) & 1) != 0)
                comDb.addBusyPort(num + 1);
        }

        LocalFree(pBuf);

    } while (res == IDTRYAGAIN);

    if (res != IDCONTINUE) {
        SetLastError(ERROR_CANCELLED);
        return FALSE;
    }

    return TRUE;
}

static bool ClaimReleasePort(DWORD num, bool claim)
{
    if (num == 0) {
        SetLastError(ERROR_INVALID_PARAMETER);
        return FALSE;
    }

    int res;

    do {
        res = IDCONTINUE;

        HCOMDB  hComDB;
        LONG err;

        err = ComDBOpen(&hComDB);

        if (err != ERROR_SUCCESS) {
            qDebug() << "ComDBOpen() failed with error:" << err;
                res = IDCANCEL;
            continue;
        }

        if (claim) {
            err = ComDBClaimPort(hComDB, num, FALSE, NULL);

            if (err != ERROR_SUCCESS) {
                ComDBClose(hComDB);

                if (err == ERROR_SHARING_VIOLATION)
                    res = IDCANCEL;
                else {
                    qDebug() << "ComDBClaimPort(COM" << (unsigned)num << ") failed with error:" << err;
                    res = IDCANCEL;
                }

                continue;
            } else {
                qDebug() << "ComDB: COM" << unsigned(num) << "- logged as \"in use\"";
            }
        } else {
            err = ComDBReleasePort(hComDB, num);

            if (err != ERROR_SUCCESS) {
                ComDBClose(hComDB);

                qDebug() << "ComDBReleasePort(COM" << (unsigned)num << ") failed with error:" << err;
            res = IDCANCEL;
                continue;
            } else {
                qDebug() << "ComDB: COM" << unsigned(num) << "- released";
            }
        }

        ComDBClose(hComDB);
    } while (res == IDRETRY);

    if (res != IDCONTINUE) {
        SetLastError(ERROR_CANCELLED);
        return FALSE;
    }

    return TRUE;
}

static bool LoadComDbLocal(PortBusyMask &comDb)
{
    comDb.clear();

    int res;

    do {
        res = IDCONTINUE;

        LONG err;
        HKEY hKey;

        err = RegOpenKeyEx(HKEY_LOCAL_MACHINE,
                           comDbLocalKey,
                           0,
                           KEY_READ,
                           &hKey);

        if (err != ERROR_SUCCESS) {
            if (err != ERROR_FILE_NOT_FOUND) {
                qDebug() << "RegOpenKeyEx(" << comDbLocalKey << ") failed with error:" << err;
                res = IDCANCEL;
            }

            continue;
        }

        DWORD bufSize = 1;

        err = RegQueryValueEx(hKey,
                              comDbLocalName,
                              NULL,
                              NULL,
                              NULL,
                              &bufSize);

        if (err != ERROR_SUCCESS) {
            RegCloseKey(hKey);

            if (err != ERROR_FILE_NOT_FOUND) {
                qDebug() << "RegQueryValueEx(" << comDbLocalKey << "\\" << comDbLocalName << ") failed with error:" << err;
                res = IDCANCEL;
            }

            continue;
        }

        BYTE *pBuf = (BYTE *)LocalAlloc(LPTR, bufSize);

        if (!pBuf) {
            RegCloseKey(hKey);

            qDebug() << "LocalAlloc(" << (unsigned long)bufSize << ") failed with ERROR_NOT_ENOUGH_MEMORY";
            res = IDCANCEL;
            continue;
        }

        err = RegQueryValueEx(hKey,
                              comDbLocalName,
                              NULL,
                              NULL,
                              pBuf,
                              &bufSize);

        RegCloseKey(hKey);

        if (err != ERROR_SUCCESS) {
            LocalFree(pBuf);

            if (err != ERROR_FILE_NOT_FOUND) {
                qDebug() << "RegQueryValueEx(" << comDbLocalKey << "\\" << comDbLocalName << ") failed with error:" << err;
                res = IDCANCEL;
            }

            continue;
        }

        DWORD maxPortsReported = bufSize*8;

        for (DWORD num = 0 ; num < maxPortsReported ; num++) {
            if (((pBuf[num/8] >> (num%8)) & 1) != 0)
                comDb.addBusyPort(num + 1);
        }

        LocalFree(pBuf);

    } while (res == IDTRYAGAIN);

    if (res != IDCONTINUE) {
        SetLastError(ERROR_CANCELLED);
        return FALSE;
    }

    return TRUE;
}

static bool AddComNames(
    HDEVINFO hDevInfo,
    PSP_DEVINFO_DATA pDevInfoData,
    ConstDevicePropertiesPtr /*pDevProperties*/,
    bool * /*pRebootRequired*/,
    void *pParam)
{
    int i = PortNum::GetPortNum(hDevInfo, pDevInfoData);

    if (i < 0)
        return TRUE;

    for (int j = 0 ; j < 2 ; j++) {
        char phPortName[20];

        snprintf(phPortName, sizeof(phPortName)/sizeof(phPortName[0]), "%s%d",
                 j ? C0C_PREF_PORT_NAME_B : C0C_PREF_PORT_NAME_A, i);

        PortParameters portParameters(C0C_SERVICE, phPortName);

        if (!portParameters.load())
            return FALSE;

        char portName[20];

        QString portNameStr = portParameters.getPortName();
        strncpy_s(portName, portNameStr.toLocal8Bit().constData(), sizeof(portName) - 1);
        portName[sizeof(portName) - 1] = '\0';

        WORD num = name2num(portName);

        if (num > 0)
            ((PortBusyMask *)pParam)->addBusyPort(num);
    }

    return TRUE;
}

// extern int EnumDevices(
//     DeviceEnumFilter pFilter,
//     ConstDevicePropertiesPtr pDevProperties,
//     bool *pRebootRequired,
//     DeviceCallback pDevCallBack,
//     void *pDevCallBackParam);

static bool LoadComNames(DeviceEnumFilter pFilter, PortBusyMask &comDb)
{
    comDb.clear();

    DeviceProperties devProperties;

    devProperties.setDeviceId(C0C_BUS_DEVICE_ID);

    if (EnumDevices(pFilter, &devProperties, nullptr, AddComNames, &comDb) < 0)
        return FALSE;

    return TRUE;
}

static bool SaveComDbLocal(const PortBusyMask &comDb)
{
    int res;

    do {
        res = IDCONTINUE;

        LONG err;
        HKEY hKey;

        err = RegCreateKeyEx(HKEY_LOCAL_MACHINE,
                             comDbLocalKey,
                             0,
                             NULL,
                             0,
                             KEY_WRITE,
                             NULL,
                             &hKey,
                             NULL);

        if (err != ERROR_SUCCESS) {
            qDebug() << "RegCreateKeyEx(" << comDbLocalKey << ") failed with error:" << err;
            res = IDCANCEL;
            continue;
        }

        DWORD maxPortsReported = 0;

        for (DWORD num = 0 ; num < COMDB_MAX_PORTS_ARBITRATED ; num++) {
            if (comDb.isPortBusy(num + 1))
                maxPortsReported = num + 1;
        }

        if (!maxPortsReported) {
            err = RegDeleteValue(hKey, comDbLocalName);

            RegCloseKey(hKey);

            if (err != ERROR_SUCCESS && err != ERROR_FILE_NOT_FOUND) {
                qDebug() << "RegDeleteValue(" << comDbLocalKey << "\\" << comDbLocalName << ") failed with error:" << err;
                res = IDCANCEL;
            }

            continue;
        }

        DWORD bufSize = (maxPortsReported + 7)/8;
        BYTE *pBuf = (BYTE *)LocalAlloc(LPTR, bufSize);

        if (!pBuf) {
            RegCloseKey(hKey);

            qDebug() << "LocalAlloc(" << (unsigned long)bufSize << ") failed with ERROR_NOT_ENOUGH_MEMORY";
            res = IDCANCEL;
            continue;
        }

        for (DWORD num = 0 ; num < maxPortsReported ; num++) {
            if (comDb.isPortBusy(num + 1))
                pBuf[num/8] |= (1 << (num%8));
        }

        err = RegSetValueEx(hKey,
                            comDbLocalName,
                            NULL,
                            REG_BINARY,
                            pBuf,
                            bufSize);

        RegCloseKey(hKey);
        LocalFree(pBuf);

        if (err != ERROR_SUCCESS) {
            qDebug() << "RegSetValueEx(" << comDbLocalKey << "\\" << comDbLocalName << ") failed with error:" << err;
            res = IDCANCEL;
            continue;
        }
    } while (res == IDTRYAGAIN);

    if (res != IDCONTINUE) {
        SetLastError(ERROR_CANCELLED);
        return FALSE;
    }

    return TRUE;
}

bool ComDbGetInUse(const char *pPortName, bool &inUse)
{
    WORD num = name2num(pPortName);

    if (num == 0) {
        inUse = FALSE;  // not arbitered by ComDB
    } else {
        PortBusyMask comDb;

        if (!LoadComDb(comDb))
            return FALSE;

        inUse = comDb.isPortBusy(num);
    }

    return TRUE;
}

void ComDbSync(DeviceEnumFilter pFilter)
{
    PortBusyMask comNames;

    if (!LoadComNames(pFilter, comNames))
        return;

    PortBusyMask comDbLocal;

    if (!LoadComDbLocal(comDbLocal))
        return;

    for (DWORD num = 0 ; num < COMDB_MAX_PORTS_ARBITRATED ; num++) {
        if (!comNames.isPortBusy(num + 1)) {
            if (comDbLocal.isPortBusy(num + 1)) {
                if (ClaimReleasePort(num + 1, FALSE)) {
                    comDbLocal.removeBusyPort(num + 1);
                    SaveComDbLocal(comDbLocal);
                }
            }
        } else {
            if (!comDbLocal.isPortBusy(num + 1)) {
                if (ClaimReleasePort(num + 1, TRUE)) {
                    comDbLocal.addBusyPort(num + 1);
                    SaveComDbLocal(comDbLocal);
                }
            }
        }
    }
}

DWORD ComDbQueryNames(char *pBuf, DWORD maxChars)
{
    PortBusyMask comDb;

    if (!LoadComDb(comDb))
        return 0;

    char *pB = pBuf;
    for (DWORD num = 0 ; num < COMDB_MAX_PORTS_ARBITRATED ; num++) {
        if (comDb.isPortFree(num))
            continue;
        int len = snprintf(pB, maxChars, "COM%u", (unsigned)(num + 1));
        if (len < 0) {
            SetLastError(ERROR_INSUFFICIENT_BUFFER);
            return 0;
        }
        pB += len + 1;
        maxChars -= len + 1;
    }
    if (pB == pBuf) {
        SetLastError(ERROR_FILE_NOT_FOUND);
        return 0;
    }
    if (maxChars == 0) {
        SetLastError(ERROR_INSUFFICIENT_BUFFER);
        return 0;
    }
    *pB++ = 0;
    return (DWORD)(pB - pBuf);
}

bool ComDbClaim(const char *pPortName)
{
    return ClaimReleasePort(name2num(pPortName), TRUE);
}

bool ComDbRelease(const char *pPortName)
{
    return ClaimReleasePort(name2num(pPortName), FALSE);
}

bool ComDbIsValidName(const char *pPortName)
{
    if (name2num(pPortName) == 0) {
        qDebug() << "The port name" << pPortName << "is not in COM<1-" << (unsigned)COMDB_MAX_PORTS_ARBITRATED << "> format.";

        return FALSE;
    }

    return name2num(pPortName) != 0;
}
